/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2019-2022. All rights reserved.
 * Description: 集合通信hcom头文件
 * Author: wangyahui
 * Create: 2019-11-28
 */

#ifndef H_COM_PUB_H
#define H_COM_PUB_H

#include <map>
#include <mutex>

#include "hccl_comm_pub.h"
#include "hcom_common.h"

#include "remote_access.h"
#include "es_private.h"

enum class CommNumHcom {
    COMM_VALUE_DEFAULT = 0, // 默认值为图模式
    COMM_VALUE_RESERVED
};

/* hccl算子类型 */
const std::string HCCL_KERNEL_OP_TYPE_BROADCAST = "HcomBroadcast";
const std::string HCCL_KERNEL_OP_TYPE_SCATTER = "HcomScatter";
const std::string HCCL_KERNEL_OP_TYPE_ALLREDUCE = "HcomAllReduce";
const std::string HCCL_KERNEL_OP_TYPE_ALLGATHER = "HcomAllGather";
const std::string HCCL_KERNEL_OP_TYPE_ALLGATHERV = "HcomAllGatherV";
const std::string HCCL_KERNEL_OP_TYPE_REDUCESCATTER = "HcomReduceScatter";
const std::string HCCL_KERNEL_OP_TYPE_SEND = "HcomSend";
const std::string HCCL_KERNEL_OP_TYPE_RECEIVE = "HcomReceive";
const std::string HCCL_KERNEL_OP_TYPE_REDUCE = "HcomReduce";
const std::string HCCL_KERNEL_OP_TYPE_REMOTE_READ = "HcomRemoteRead";
const std::string HCCL_KERNEL_OP_TYPE_REMOTE_WRITE = "HcomRemoteWrite";
const std::string HCCL_KERNEL_OP_TYPE_REMOTE_REF_READ = "HcomRemoteRefRead";
const std::string HCCL_KERNEL_OP_TYPE_REMOTE_SCATTER_WRITE = "HcomRemoteScatterWrite";
const std::string HCCL_KERNEL_OP_TYPE_ALLTOALLV = "HcomAllToAllV";
const std::string HCCL_KERNEL_OP_TYPE_ALLTOALLVC = "HcomAllToAllVC";
const std::string HCCL_KERNEL_OP_TYPE_GATHER_ALLTOALLV = "HcomGatherAllToAllV";
const std::string HCCL_KERNEL_OP_TYPE_ALLTOALL = "HcomAllToAll";
const std::string HCCL_KERNEL_OP_TYPE_REDUCESCATTERV = "HcomReduceScatterV";

HcclResult HcomRpcGetDeviceType(const char *group, DevType &deviceType);
HcclResult HcomRpcGetHccsLinkNum(const char *group, u32 &numHccsLink);
HcclResult HcomRpcGetRankSize(const char *group, u32 *rankSize);
HcclResult HcomRpcGetInitStatus(bool& initiated);
bool HcomIsRpcComm(const char *group);
bool HcomIsNormalComm(const char *group);
HcclResult HcomRpcGetRankId(const char *group, u32 *rankId);
HcclResult HcomRpcUnloadTask(const std::string group, const char *tag);
HcclResult HcomRpcSetWorkspaceResource(const std::string &tag, const char *group, std::vector<rtStream_t> stream,
    void *memPtr, u64 maxSize);
HcclResult HcomRpcGetTopoDesc(const char *group, HcclTopoDescs *topoDescs, uint32_t topoSize);

// 生成通信域标识符
HcclResult HcomGenerateCommId(hccl::HcclCommParams &params);
HcclResult HcomSetProfilingMode(HcomProfilingMode profilingMode, std::string profilingOption);
HcclResult HcomCheckCommValidity(const char* group);
HcclResult HcomGetCCLBufferAvailableSize(u64 &size);
HcclResult HcomGetRemoteAccess(std::shared_ptr<hccl::RemoteAccess> &remoteAccess);
HcclResult HcomGetDevId(const char *group, s32 *devId);
HcclResult HcclCommGraphGetDevId(s64 opBaseHcom, s32 *devId);
HcclResult HcclCommGraphGetIdentifier(s64 opBaseHcom, std::string &identifier);
HcclResult HcclCommGraphGetIdentifierCpu(s64 opBaseHcom, std::string &identifier);
HcclResult HcomInitByMasterInfo(const std::string &masterIp, const std::string &masterPort,
    const std::string &masterDeviceId, const std::string &rankSize, const std::string &rankIp);
DevType HcomGetDeviceType();
HcclResult CallMsprofReportHostApi(hccl::hcclComm* hcclComm, HcclCMDType cmdType, uint64_t beginTime, u64 count,
    HcclDataType dataType);
#ifdef __cplusplus
extern "C" {
#endif
HcclResult HcomInit(const char *rankTableM, const char *identify,
    WorkMode commWorkMode = WorkMode::HCCL_MODE_NORMAL);
HcclResult HcomInitByString(const char *rankTableM, const char *identify,
    WorkMode commWorkMode = WorkMode::HCCL_MODE_NORMAL);

HcclResult HcomDestroy(void);
HcclResult HcomGetInitStatus(bool& initiated);
HcclResult HcomAllGather(const char *tag, void *inputPtr, void *outputPtr, u64 inputCount,
    HcclDataType dataType, const char *group, rtStream_t stream);
HcclResult HcomAllGatherV(const char *tag, const void *sendBuf, u64 sendCount, const void *recvBuf,
    const void *recvCounts, const void *rdispls, HcclDataType dataType, const char *group, rtStream_t stream);
HcclResult HcomAllReduce(const char *tag, void *inputPtr, void *outputPtr, u64 count,
    HcclDataType dataType, HcclReduceOp op, const char *group, rtStream_t stream);
HcclResult HcomReduce(const char *tag, void *inputPtr, void *outputPtr, u64 count, HcclDataType dataType,
    HcclReduceOp op, u32 root, const char *group, rtStream_t stream);
HcclResult HcomBroadcast(const char *tag, void *ptr, u64 count, HcclDataType dataType, u32 root,
    const char *group, rtStream_t stream);
HcclResult HcomReduceScatter(const char *tag, void *inputPtr, void *outputPtr, u64 count,
    HcclDataType dataType, HcclReduceOp op, const char *group, rtStream_t stream);
HcclResult HcomReduceScatterV(const char *tag, void *sendBuf, const void *sendCounts, const void *sdispls,
    void *recvBuf, u64 recvCount, HcclDataType dataType, HcclReduceOp op, const char *group, rtStream_t stream);
HcclResult HcomSend(const char *tag, void *inputPtr, u64 count, HcclDataType dataType,
    u32 destRank, u32 srTag, const char *group, rtStream_t stream);
HcclResult HcomReceive(const char *tag, void *outputPtr, u64 count, HcclDataType dataType,
    u32 srcRank, u32 srTag, const char *group, rtStream_t stream);

// pytorch单算子通信域复用 start
HcclResult HcclCommGraphAllGather(const char *tag, void *inputPtr, void *outputPtr, u64 inputCount,
    HcclDataType dataType, s64 opBaseHcom, rtStream_t stream);
HcclResult HcclCommGraphAllReduce(const char *tag, void *inputPtr, void *outputPtr, u64 count,
    HcclDataType dataType, HcclReduceOp op, s64 opBaseHcom, rtStream_t stream);
HcclResult HcclCommGraphReduce(const char *tag, void *inputPtr, void *outputPtr, u64 count, HcclDataType dataType,
    HcclReduceOp op, u32 root, s64 opBaseHcom, rtStream_t stream);
HcclResult HcclCommGraphBroadcast(const char *tag, void *ptr, u64 count, HcclDataType dataType, u32 root,
    s64 opBaseHcom, rtStream_t stream);
HcclResult HcclCommGraphReduceScatter(const char *tag, void *inputPtr, void *outputPtr, u64 count,
    HcclDataType dataType, HcclReduceOp op, s64 opBaseHcom, rtStream_t stream);
HcclResult HcclCommGraphSend(const char *tag, void *inputPtr, u64 count, HcclDataType dataType,
    u32 destRank, u32 srTag, s64 opBaseHcom, rtStream_t stream);
HcclResult HcclCommGraphReceive(const char *tag, void *outputPtr, u64 count, HcclDataType dataType,
    u32 srcRank, u32 srTag, s64 opBaseHcom, rtStream_t stream);
HcclResult HcclCommGraphAlltoAllV(const void *sendBuf, const void *sendCounts, const void *sdispls,
    HcclDataType sendType, const void *recvBuf, const void *recvCounts, const void *rdispls,
    HcclDataType recvType, s64 opBaseHcom, rtStream_t stream, const char *tag);
HcclResult HcclCommGraphAlltoAllVC(const void *sendBuf, const void *sendCountMatrix, HcclDataType sendType,
    const void *recvBuf, HcclDataType recvType, s64 opBaseHcom, rtStream_t stream, const char *tag);
HcclResult HcclCommGraphGetAlltoAllStagedWorkSpaceMemSize(s64 opBaseHcom, u64 *sendCounts, u64 *sdispls,
    HcclDataType sendType, u64 *recvCounts, u64 *rdispls, HcclDataType recvType, u64 &memSize);
HcclResult HcclCommGraphSetWorkspaceResource(const std::string &tag, s64 opBaseHcom, std::vector<rtStream_t> stream,
    void *memPtr, u64 maxSize);
HcclResult HcclCommGraphGetAllReduceScratchSize(s64 opBaseHcom, const u32 count, const HcclDataType dataType,
    u64 &outScratchSize);
HcclResult HcclCommGraphGetRankSize(s64 opBaseHcom, u32 *rankSize);
HcclResult HcclCommGraphGetRankId(s64 opBaseHcom, u32 *rankId);
HcclResult HcclCommGraphGetWorkspaceSubStreamNum(s64 opBaseHcom, u64 &streamNum,
    u64 dataSize = 0, HcclCMDType opType = HcclCMDType::HCCL_CMD_INVALID);
HcclResult HcclCommGraphSetIsPytorchComm();
// pytorch单算子通信域复用 end

HcclResult HcomGetSplitStrategy(const char *group, const struct model_feature *feature,
    std::vector<u32>& segmentIdx, bool &configured, GradSplitForceMode force = GradSplitForceMode::FORCE_NONE,
    OriginalGraphShapeType shapeType = OriginalGraphShapeType::KNOWN_SHAPE);
HcclResult HcomGetWorkspaceSubStreamNum(const char *group, u64 &streamNum, u64 dataSize = 0,
    HcclCMDType optype = HcclCMDType::HCCL_CMD_INVALID);
HcclResult HcomGetWorkspaceMemSize(const std::string &opType, u64 count,
    HcclDataType dataType, const char *group, u64 &memSize);
HcclResult HcomGetAlltoAllStagedWorkSpaceMemSize(const char *group, u64 *sendCounts, u64 *sdispls,
    HcclDataType sendType, u64 *recvCounts, u64 *rdispls, HcclDataType recvType, u64 &memSize);
HcclResult HcomGetAlltoAllvcStagedWorkSpaceMemSize(const char *group,
    std::vector<SendRecvInfo> &allMeshAggregationSendRecvInfo, u64 &memSize);
HcclResult HcomGetAllReduceScratchSize(const char *group, const u32 count, const HcclDataType dataType,
    u64 &outScratchSize);
HcclResult HcomSetWorkspaceResource(const std::string &tag, const char *group, std::vector<rtStream_t> stream,
    void *memPtr, u64 maxSize);
HcclResult HcomSetAttachedStream(const char *group, const std::vector<rtStream_t> &stream);
HcclResult HcclCommSetAttachedStream(s64 opBaseHcom, const std::vector<rtStream_t> &stream);
void HcomSetAutoTuneMode(bool autoTuneMode);
HcclResult HcomSetExecTimeOut(const char *execTimeOut);
HcclResult HcomSetAlgorithm(const char* algo);
HcclResult HcomGetAlgorithm(u32 level, std::string &algo);
HcclResult HcomSetDeterministic(bool isDeterministic);
HcclResult HcomGetBandWidthPerNPU(u32 level, float &bandWidth);
HcclResult HcomReleaseSubComms();
HcclResult HcomAlltoAllV(const void *sendBuf, const void *sendCounts, const void *sdispls, HcclDataType sendType,
    const void *recvBuf, const void *recvCounts, const void *rdispls, HcclDataType recvType,
    const char *group, rtStream_t stream, const char *tag);
HcclResult HcomAlltoAllVC(const void *sendBuf, const void *sendCountMatrix, HcclDataType sendType,
    const void *recvBuf, HcclDataType recvType, const char *group, rtStream_t stream, const char *tag);
HcclResult HcomGetServerNumAndDeviceNumPerServer(u32 &serverNum, u32 &deviceNumPerServer, u32 &deviceNumPerAggregation);
HcclResult HcclCommGraphUnloadTask(s64 opBaseHcom, const char *tag);

HcclResult HcomSetRankTableImpl(const char *rankTableStr);
HcclResult HcomGetActualRankSizeImpl(const char *group, u32 *rankSize);
HcclResult HcomSetGlobalWorkSpace(const char *group, std::vector<void *> &globalWorkSpaceAddr);
HcclResult HcclCommSetGlobalWorkSpace(s64 opBaseHcom, std::vector<void *> &globalWorkSpaceAddr);

HcclResult HcomGetandClearOverFlowTasks(const char *group, std::vector<hccl::HcclDumpInfo> &hcclDumpInfo);
HcclResult HcclCommGetandClearOverFlowTasks(s64 opBaseHcom, std::vector<hccl::HcclDumpInfo> &hcclDumpInfo);
HcclResult HcomSupportDeterministicOptim(const char *group, bool &isDeterministicOptim);
HcclResult HcclCommSupportDeterministicOptim(s64 opBaseHcom, bool &isDeterministicOptim);
HcclResult HcomGetHccsLinkNum(const char *group, u32 &numHccsLink);
HcclResult HcomAlltoAll(const void *sendBuf, u64 sendCount, HcclDataType sendType,
                        const void *recvBuf, u64 recvCount, HcclDataType recvType,
                        const char *group, rtStream_t stream, const char *tag);
HcclResult HcomUnloadTask(const std::string group, const char *tag);
HcclResult HcomRemoteLookup(s32 tag, s32 insertFlag, void *stream, s32 &intZerocpyFlag, s32 &outZerocpyFlag,
    RemoteLookupPubParams &lookupPubInfo, u64 keyNum = 0);
HcclResult HcomSetQosCfg(const char *group, const u32 qosCfg);
HcclResult HcclCommSetQosCfg(s64 opBaseHcom, const u32 qosCfg);
HcclResult HcomResetQosCfg(const char *group);
HcclResult HcclCommResetQosCfg(s64 opBaseHcom);
HcclResult HcomCollRemoteUpdate(s32 tag, void* keys, u32 keyMaxNum, void *values, void *tableId, s32 insertOption,
    s32 valueItemSize, const char *group, void *stream, s32 maxEmbeddingDim, u64 keyNum = 0);
HcclResult HcomCollRemoteLookup(s32 tag, void* keys, u32 keyMaxNum, void *values, void *tableId, s32 insertFlag,
    s32 valueItemSize, const char *group, void *stream, s32 flags, s32 maxEmbeddingDim, u64 keyNum = 0);
HcclResult HcomGather(const char *tag, void *inputPtr, void *outputPtr, u32 rootRank, u64 inputCount,
    HcclDataType dataType, const char *group, rtStream_t stream);
HcclResult GenerateCclOpTag(const std::string &opType, int64_t &hcomComm,
    std::string& group, std::string &sTag);
HcclResult HcomCollRemoteLookupPaired(const HcomRemoteOperationParams &params);
HcclResult HcomCollRemoteLookupUniquedAndPaired(const HcomRemoteOperationParams &params);
HcclResult HcomCollRemoteUpdatePaired(const HcomRemoteOperationParams &params);
HcclResult HcclReStartVnic(s64 comm);

HcclResult HcomClearAivSyncBuf(const char *group, bool aivClearEnable);
HcclResult HcclCommGraphClearAivSyncBuf(s64 comm, bool aivClearEnable);

#ifdef __cplusplus
}
#endif
#endif  // H_COM_PUB_H
